from pathlib import Path
from typing import List, Optional
from openai import OpenAI
import os
import sys
from dotenv import load_dotenv
import logging

logger = logging.getLogger(__name__)

class QuestionGenerator:
    def __init__(self):
        load_dotenv()
        api_key = os.getenv("KIMI_API_KEY")
        if not api_key:
            raise ValueError("KIMI_API_KEY not found in environment variables")
            
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.moonshot.cn/v1"
        )
        self.history = []
        self.system_prompt = {
            "role": "system",
            "content": """你是一个专业的出题助手，擅长根据文档内容生成高质量的题目。
            你会根据用户提供的文档内容，生成不同类型的题目，包括：
            1. 选择题：包含4个选项，其中1个正确答案
            2. 填空题：需要填写关键信息
            3. 简答题：需要简要回答
            4. 论述题：需要详细论述
            
            每道题目都会包含：
            1. 题目内容

            同时在最后给出每一道题目的
            2. 参考答案
            3. 详细解析
            4. 难度等级（简单/中等/困难

            请确保题目：
            1. 符合文档内容
            2. 难度适中
            3. 答案准确
            4.参考答案和详细解析以及难度等级在最后按照题目顺序实现
            5. 解析详细"""
        }
    
    def _add_to_history(self, role: str, content: str):
        """添加消息到历史记录"""
        self.history.append({"role": role, "content": content})
    
    def _get_messages(self, user_input: str) -> List[dict]:
        """构建消息列表"""
        messages = [self.system_prompt]
        messages.extend(self.history)
        messages.append({"role": "user", "content": user_input})
        return messages
    
    def chat(self, message: str) -> str:
        """基础对话功能"""
        messages = self._get_messages(message)
        
        completion = self.client.chat.completions.create(
            model="moonshot-v1-8k",
            messages=messages,
            temperature=0.3
        )
        
        response = completion.choices[0].message.content
        self._add_to_history("user", message)
        self._add_to_history("assistant", response)
        return response
    
    def process_document(self, file_path: str, question_type: Optional[str] = None, count: int = 5, difficulty: str = "中等", context: Optional[str] = None) -> str:
        """处理文档并生成题目"""
        try:
            logger.info(f"开始处理文档: {file_path}")
            
            # 上传文档
            file_object = self.client.files.create(
                file=Path(file_path),
                purpose="file-extract"
            )
            
            logger.info(f"文件上传成功，ID: {file_object.id}")
            
            # 获取文档内容
            file_content = self.client.files.content(file_id=file_object.id).text
            
            logger.info("成功获取文件内容")
            
            # 构建题目生成要求
            requirements = []
            requirements.append(f"请根据文档内容生成{count}道{question_type if question_type else '不同类型的'}题目")
            requirements.append(f"难度为{difficulty}")
            
            # 如果用户提供了具体要求，加入到生成要求中
            if context and context.strip():
                requirements.append(f"具体要求：{context.strip()}")
            
            requirements.append("每道题目都要有详细的解析")
            
            # 构建完整的题目生成要求
            user_content = "，".join(requirements) + "。"
            
            # 构建消息列表
            messages = [
                {
                    "role": "system",
                    "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文的对话。你会为用户提供安全，有帮助，准确的回答。你是一个专业的出题助手，擅长根据文档内容生成高质量的题目。"
                },
                {
                    "role": "system",
                    "content": file_content
                },
                {
                    "role": "user",
                    "content": user_content
                }
            ]
            
            logger.info(f"文档题目生成要求: {user_content}")
            
            # 调用 chat completion
            completion = self.client.chat.completions.create(
                model="moonshot-v1-32k",
                messages=messages,
                temperature=0.3
            )
            
            logger.info("成功生成题目")
            
            # 清理上传的文件
            self.client.files.delete(file_id=file_object.id)
            logger.info(f"已清理上传的文件: {file_object.id}")
            
            return completion.choices[0].message.content
        except Exception as e:
            logger.error(f"处理文档失败: {str(e)}")
            raise
    
    def generate_questions(self, 
                         core_words: Optional[str] = None,
                         question_type: Optional[str] = None,
                         count: int = 5,
                         difficulty: str = "中等",
                         file_path: Optional[str] = None) -> str:
        """生成题目"""
        try:
            if file_path:
                return self.process_document(file_path, question_type, count, difficulty, core_words)
            else:
                # 构建题目生成要求
                requirements = []
                if core_words:
                    requirements.append(f"请根据核心词'{core_words}'生成{count}道题目")
                else:
                    requirements.append(f"请生成{count}道题目")
                
                if question_type:
                    requirements.append(f"题目类型为{question_type}")
                
                requirements.append(f"难度为{difficulty}")
                requirements.append("每道题目都要有详细的解析")
                
                # 构建完整的题目生成要求
                content = "，".join(requirements) + "。"
                
                messages = [
                    self.system_prompt,
                    {
                        "role": "user",
                        "content": content
                    }
                ]
                
                logger.info(f"生成题目要求: {content}")
                
                completion = self.client.chat.completions.create(
                    model="moonshot-v1-8k",
                    messages=messages,
                    temperature=0.3
                )
                
                response = completion.choices[0].message.content
                logger.info("题目生成完成")
                
                return response
                
        except Exception as e:
            logger.error(f"生成题目时出错：{str(e)}")
            return f"生成题目时出错：{str(e)}"
    
    def clear_history(self):
        """清空历史记录"""
        self.history = [] 

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("请提供命令: chat/process_document/generate_questions")
        sys.exit(1)
        
    try:
        generator = QuestionGenerator()
        command = sys.argv[1]
        
        if command == "chat":
            if len(sys.argv) < 3:
                print("请提供聊天消息")
                sys.exit(1)
            response = generator.chat(sys.argv[2])
            print(response)
            
        elif command == "process_document":
            if len(sys.argv) < 3:
                print("请提供文档路径")
                sys.exit(1)
            response = generator.process_document(sys.argv[2])
            print(response)
            
        elif command == "generate_questions":
            file_path = sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] else None
            question_type = sys.argv[3] if len(sys.argv) > 3 and sys.argv[3] else None
            count = int(sys.argv[4]) if len(sys.argv) > 4 else 5
            difficulty = sys.argv[5] if len(sys.argv) > 5 else "中等"
            context = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else None
            
            response = generator.generate_questions(
                file_path=file_path,
                question_type=question_type,
                count=count,
                difficulty=difficulty,
                context=context
            )
            print(response)
            
        else:
            print(f"未知命令: {command}")
            sys.exit(1)
            
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)
